import torch


def rnncell():
    # how to use RNNCell
    batch_size = 1
    seq_len = 3
    input_size = 4
    hidden_size = 2
    cell = torch.nn.RNNCell(input_size=input_size, hidden_size=hidden_size)
    # (seq, batch, features)
    dataset = torch.randn(seq_len, batch_size, input_size)
    hidden = torch.zeros(batch_size, hidden_size)
    for idx, input in enumerate(dataset):
        print('=' * 20, idx, '=' * 20)
        print('Input size: ', input.shape)
        hidden = cell(input, hidden)
        print('outputs size: ', hidden.shape)
        print(hidden)


def rnn():
    # how to use RNN
    batch_size = 1
    seq_len = 3
    input_size = 4
    hidden_size = 2
    num_layers = 1
    cell = torch.nn.RNN(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers)
    # (seqLen, batchSize, inputSize)
    inputs = torch.randn(seq_len, batch_size, input_size)
    hidden = torch.zeros(num_layers, batch_size, hidden_size)
    out, hidden = cell(inputs, hidden)
    print('Output size:', out.shape)
    print('Output:', out)
    print('Hidden size: ', hidden.shape)
    print('Hidden: ', hidden)


# Design Model
class Model(torch.nn.Module):
    def __init__(self, input_size, hidden_size, batch_size):
        super(Model, self).__init__()
        self.batch_size = batch_size
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.rnncell = torch.nn.RNNCell(input_size=self.input_size,
                                        hidden_size=self.hidden_size)

    def forward(self, input, hidden):
        hidden = self.rnncell(input, hidden)
        return hidden

    def init_hidden(self):
        return torch.zeros(self.batch_size, self.hidden_size)


def rnnModel():
    input_size = 4
    hidden_size = 4
    batch_size = 1
    idx2char = ['e', 'h', 'l', 'o']
    x_data = [1, 0, 2, 2, 3]  # hello
    y_data = [0, 2, 2, 3, 1]  # elloh
    one_hot_lookup = [[1, 0, 0, 0],
                      [0, 1, 0, 0],
                      [0, 0, 1, 0],
                      [0, 0, 0, 1]]
    x_one_hot = [one_hot_lookup[x] for x in x_data]
    inputs = torch.Tensor(x_one_hot).view(-1, batch_size, input_size)
    labels = torch.LongTensor(y_data).view(-1, 1)

    net = Model(input_size, hidden_size, batch_size)

    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(net.parameters(), lr=0.1)

    for epoch in range(30):
        loss = 0
        optimizer.zero_grad()
        hidden = net.init_hidden()
        print('Predicted string: ', end='')
        for input, label in zip(inputs, labels):
            hidden = net(input, hidden)
            loss += criterion(hidden, label)
            _, idx = hidden.max(dim=1)
            print(idx2char[idx.item()], end='')
        loss.backward()
        optimizer.step()
        print(', Epoch [%d/15] loss=%.4f' % (epoch + 1, loss.item()))


if __name__ == '__main__':
    # rnncell()
    # rnn()
    rnnModel()
    print("====over=====")
